{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ddc557b7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "done\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import wfdb\n",
    "import ast\n",
    "\n",
    "def load_raw_data(df, sampling_rate, path):\n",
    "    if sampling_rate == 100:\n",
    "        data = [wfdb.rdsamp(path+f) for f in df.filename_lr]\n",
    "    else:\n",
    "        data = [wfdb.rdsamp(path+f) for f in df.filename_hr]\n",
    "    data = np.array([signal for signal, meta in data])\n",
    "    return data\n",
    "\n",
    "path = 'C:/Users/mstew/OneDrive/Curtin Univeristy/COMP6011/Task 3/ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.1/'\n",
    "sampling_rate=100\n",
    "calssificatin_type = \"superclasses\"    #{\"binary\",\"superclasses\",\"subclasses\"}\n",
    "\n",
    "lead_types={\"lead-I\":[1,2,3,4,5,6,7,8,9,10,11], \"bipolar-limb\":[3,4,5,6,7,8,9,10,11] , \"unipolar-limb\":[0,1,2,6,7,8,9,10,11], \"limb-leads\":[6,7,8,9,10,11] , \"precordial-leads\":[0,1,2,3,4,5],\"all-lead\":[]}\n",
    "lead_name=\"all-lead\"\n",
    "\n",
    "# load and convert annotation data\n",
    "Y = pd.read_csv(path+'ptbxl_database.csv', index_col='ecg_id')\n",
    "Y.scp_codes = Y.scp_codes.apply(lambda x: ast.literal_eval(x))\n",
    "\n",
    "#\n",
    "# Load raw signal data\n",
    "X = load_raw_data(Y, sampling_rate, path)\n",
    "\n",
    "# Load scp_statements.csv for diagnostic aggregation\n",
    "agg_df = pd.read_csv(path+'scp_statements.csv', index_col=0)\n",
    "agg_df = agg_df[agg_df.diagnostic == 1]\n",
    "\n",
    "#\n",
    "def aggregate_superclass_diagnostic(y_dic):\n",
    "    tmp = []\n",
    "    for key in y_dic.keys():\n",
    "        if key in agg_df.index:\n",
    "            tmp.append(agg_df.loc[key].diagnostic_class)\n",
    "    return list(set(tmp))\n",
    "\n",
    "\n",
    "def aggregate_subclass_diagnostic(y_dic):\n",
    "    tmp = []\n",
    "    for key in y_dic.keys():\n",
    "        if key in agg_df.index:\n",
    "            tmp.append(agg_df.loc[key].diagnostic_subclass)\n",
    "    ret = list(set(tmp))\n",
    "    return ret\n",
    "\n",
    "\n",
    "if  calssificatin_type==\"superclasses\":\n",
    "    Y['diagnostic_superclass'] = Y.scp_codes.apply(aggregate_subclass_diagnostic)\n",
    "else:\n",
    "    Y['diagnostic_superclass'] = Y.scp_codes.apply(aggregate_superclass_diagnostic)\n",
    "\n",
    "\n",
    "\n",
    "# Split data into train and test\n",
    "test_fold = 10\n",
    "# Train\n",
    "X_train = X[np.where(Y.strat_fold != test_fold)]\n",
    "y_train = Y[(Y.strat_fold != test_fold)].diagnostic_superclass\n",
    "\n",
    "\n",
    "\n",
    "# Test\n",
    "X_test = X[np.where(Y.strat_fold == test_fold)]\n",
    "y_test = Y[Y.strat_fold == test_fold].diagnostic_superclass\n",
    "\n",
    "y_train= y_train.tolist()\n",
    "y_test= y_test.tolist()\n",
    "\n",
    "if  calssificatin_type==\"binary\":\n",
    "    count =0\n",
    "    for i in y_train:\n",
    "        if 'MI' in i or 'HYP' in i or 'CD' in i or 'STTC' in i:\n",
    "            y_train[count] = 0\n",
    "        elif 'NORM' in i:\n",
    "            y_train[count] = 1\n",
    "\n",
    "        else:\n",
    "            y_train[count] = 0\n",
    "\n",
    "\n",
    "        count = count+1\n",
    "\n",
    "\n",
    "    count =0\n",
    "    for i in y_test:\n",
    "        if 'MI' in i or 'HYP' in i or 'CD' in i or 'STTC' in i:\n",
    "            y_test[count] = 0\n",
    "        elif 'NORM' in i:\n",
    "            y_test[count] = 1\n",
    "        else:\n",
    "            y_test[count] = 0\n",
    "        count = count + 1\n",
    "  \n",
    "\n",
    "print(\"done\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1b43971d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "def mkdir_ifnotexists(dir):\n",
    "    if os.path.exists(dir):\n",
    "        return\n",
    "    os.mkdir(dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "42f542b3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI'],\n",
       " ['IRBBB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " ['NST_'],\n",
       " ['LAFB/LPFB', 'AMI', 'ISC_', 'LVH'],\n",
       " ['ISCA'],\n",
       " ['LAFB/LPFB', 'CRBBB'],\n",
       " ['IMI'],\n",
       " ['CLBBB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['NORM', '_AVB'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH', 'AMI', 'ISC_'],\n",
       " ['LVH', 'AMI', 'LAO/LAE', 'ISC_'],\n",
       " ['LVH', 'LAO/LAE', 'ISC_'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['LVH', 'NST_', 'ISC_'],\n",
       " ['LAO/LAE', 'ISC_', 'LVH', 'NST_', '_AVB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " ['LAFB/LPFB'],\n",
       " [],\n",
       " ['NST_'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['AMI'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " [],\n",
       " ['LVH', 'ISC_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM', 'IRBBB'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['CRBBB', 'AMI'],\n",
       " ['IMI'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['AMI'],\n",
       " ['STTC', 'IVCD'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['NORM'],\n",
       " ['ISCA', 'AMI', 'STTC'],\n",
       " ['LAFB/LPFB', 'IRBBB', 'NST_'],\n",
       " ['NORM'],\n",
       " ['STTC', 'LAO/LAE'],\n",
       " ['AMI', 'IVCD', 'NST_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['CLBBB'],\n",
       " ['NORM'],\n",
       " ['CLBBB'],\n",
       " ['NORM'],\n",
       " ['IVCD'],\n",
       " ['ILBBB', '_AVB', 'IMI'],\n",
       " ['NORM'],\n",
       " ['IVCD', '_AVB', 'ISCI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['AMI', 'ISCI'],\n",
       " ['NORM', 'LAFB/LPFB', 'IRBBB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NST_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['IRBBB'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['ISCA', 'AMI'],\n",
       " ['NST_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM', 'IVCD'],\n",
       " ['NORM'],\n",
       " ['AMI', 'LAO/LAE', 'RAO/RAE'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['IRBBB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI', 'IRBBB', '_AVB'],\n",
       " ['STTC'],\n",
       " ['AMI', 'IMI'],\n",
       " ['LAFB/LPFB', 'AMI', 'IRBBB'],\n",
       " ['STTC'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['AMI', '_AVB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['CLBBB', 'AMI', '_AVB', 'IMI'],\n",
       " ['NORM'],\n",
       " ['STTC', 'NST_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['CRBBB', 'NST_', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['IRBBB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'CRBBB', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['ISCI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['IRBBB', 'IMI'],\n",
       " ['LVH'],\n",
       " ['NORM', 'LAFB/LPFB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH', 'AMI', 'ISC_', 'IMI'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['LVH'],\n",
       " ['LAFB/LPFB', 'AMI', 'IMI'],\n",
       " ['ISCI'],\n",
       " ['AMI'],\n",
       " ['IVCD'],\n",
       " ['CRBBB', 'ISCI'],\n",
       " ['LVH', 'ISCI'],\n",
       " ['ISCA'],\n",
       " ['AMI', 'LAO/LAE'],\n",
       " ['ISC_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH', 'STTC', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH'],\n",
       " ['LVH'],\n",
       " ['LAO/LAE', 'IMI'],\n",
       " ['CLBBB'],\n",
       " ['IVCD', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " [],\n",
       " ['STTC'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['LAO/LAE'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['LMI'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['RAO/RAE'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH', 'LAO/LAE', 'ISC_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI'],\n",
       " ['NST_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'ISCI'],\n",
       " ['ISCA', 'IMI'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['LVH', 'AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NST_'],\n",
       " ['STTC', 'IRBBB', 'IMI'],\n",
       " ['LAFB/LPFB', 'AMI', 'IRBBB'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['ISCA', 'AMI'],\n",
       " ['AMI', 'LMI', 'ISCI'],\n",
       " ['ISC_'],\n",
       " ['NORM'],\n",
       " ['ISCA'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['LVH', 'AMI', 'ISC_'],\n",
       " ['STTC'],\n",
       " ['LVH'],\n",
       " ['NORM'],\n",
       " ['_AVB'],\n",
       " ['STTC'],\n",
       " ['LAFB/LPFB'],\n",
       " ['NORM'],\n",
       " ['ISCA'],\n",
       " ['STTC'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['LVH', 'STTC'],\n",
       " ['AMI', 'STTC', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI', 'IMI'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'STTC'],\n",
       " ['IRBBB'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['LVH', 'STTC', 'ISC_'],\n",
       " ['LVH', 'AMI', 'ISC_'],\n",
       " ['AMI', 'LAO/LAE', 'IRBBB'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'AMI'],\n",
       " ['NORM'],\n",
       " ['NORM', 'IRBBB'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " ['IRBBB'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['STTC'],\n",
       " ['LAFB/LPFB', 'RVH', 'IRBBB'],\n",
       " ['NORM'],\n",
       " ['RVH'],\n",
       " ['IVCD'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['CRBBB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH'],\n",
       " ['NORM'],\n",
       " ['AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI', 'ISCI'],\n",
       " ['NORM', 'IRBBB'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['LMI'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " [],\n",
       " [],\n",
       " ['NORM'],\n",
       " ['AMI', 'STTC'],\n",
       " ['LVH', 'NST_'],\n",
       " ['IMI'],\n",
       " ['_AVB', 'IMI'],\n",
       " ['IRBBB'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['ISC_'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " ['STTC'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM', 'IRBBB'],\n",
       " ['AMI', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH', 'AMI', 'ISC_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NST_'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM', 'IRBBB'],\n",
       " ['NORM', 'IVCD'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['ISCA'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM', 'IRBBB'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['ISCA', 'ISCI'],\n",
       " ['LVH'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " ['STTC', 'IMI'],\n",
       " ['LVH', 'NST_'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['AMI', 'IMI'],\n",
       " ['AMI', 'IVCD'],\n",
       " ['NORM', 'IVCD'],\n",
       " ['AMI', 'STTC', 'NST_'],\n",
       " ['LAO/LAE', 'ISC_', 'LVH', 'IVCD', 'IMI'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['STTC'],\n",
       " ['NORM', 'IRBBB'],\n",
       " ['AMI', 'IMI'],\n",
       " ['NST_'],\n",
       " ['NORM'],\n",
       " ['STTC', 'NST_'],\n",
       " ['IMI'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " ['STTC'],\n",
       " ['LVH'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " ['LVH', 'AMI', 'ISC_'],\n",
       " ['CLBBB', '_AVB'],\n",
       " ['CLBBB', 'ISCA'],\n",
       " ['NORM', 'IVCD', 'LAO/LAE', '_AVB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI'],\n",
       " ['NORM'],\n",
       " ['IVCD', 'LMI'],\n",
       " ['AMI', 'IRBBB'],\n",
       " ['IMI'],\n",
       " ['CRBBB', 'AMI', 'IMI'],\n",
       " ['NORM'],\n",
       " ['ISCA'],\n",
       " ['IMI'],\n",
       " ['LVH'],\n",
       " ['LVH', 'AMI'],\n",
       " ['IMI'],\n",
       " ['NST_'],\n",
       " ['IVCD', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI', 'IRBBB'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM', 'IVCD'],\n",
       " ['LAFB/LPFB', 'IRBBB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH', 'IVCD'],\n",
       " ['STTC', 'LMI'],\n",
       " ['NST_', 'IMI'],\n",
       " ['IMI'],\n",
       " ['LAFB/LPFB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['IRBBB'],\n",
       " ['NORM'],\n",
       " [],\n",
       " ['NORM', 'IRBBB'],\n",
       " ['AMI', 'IVCD', '_AVB'],\n",
       " ['NORM'],\n",
       " ['LVH'],\n",
       " ['IRBBB'],\n",
       " [],\n",
       " ['AMI', 'LMI'],\n",
       " ['ISC_', 'LVH', 'AMI', 'LAFB/LPFB', 'IMI'],\n",
       " ['LVH', 'NST_'],\n",
       " ['AMI', '_AVB'],\n",
       " ['NORM'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['NORM', 'IRBBB'],\n",
       " ['AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['NORM'],\n",
       " ['AMI'],\n",
       " ['NORM'],\n",
       " ['LVH', 'ISC_', 'IMI'],\n",
       " ['AMI'],\n",
       " ['AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'LVH'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'CRBBB'],\n",
       " ['_AVB'],\n",
       " ['AMI', 'IMI'],\n",
       " ['AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " [],\n",
       " ['CLBBB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " ['ISCA', 'AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI'],\n",
       " ['NST_'],\n",
       " ['IMI'],\n",
       " ['NORM', 'IVCD', 'LAO/LAE'],\n",
       " ['AMI'],\n",
       " ['AMI'],\n",
       " ['IMI'],\n",
       " ['LAFB/LPFB'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['LAFB/LPFB'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['IRBBB'],\n",
       " ['NORM'],\n",
       " ['NORM', 'IRBBB'],\n",
       " ['IRBBB'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI'],\n",
       " ['LVH'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['LAFB/LPFB', 'LVH', 'STTC'],\n",
       " ['NORM', 'NST_'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['LAO/LAE', 'IRBBB'],\n",
       " ['IMI'],\n",
       " ['ISC_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['CLBBB'],\n",
       " ['LAFB/LPFB'],\n",
       " ['AMI', 'IMI'],\n",
       " ['LVH'],\n",
       " ['STTC'],\n",
       " ['ISCI'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['RAO/RAE', 'IRBBB'],\n",
       " ['LVH'],\n",
       " ['NORM'],\n",
       " ['CLBBB'],\n",
       " ['NORM'],\n",
       " ['AMI'],\n",
       " ['LAFB/LPFB', 'RVH', 'IRBBB'],\n",
       " ['NORM'],\n",
       " ['LVH', 'AMI', 'ISC_'],\n",
       " ['AMI'],\n",
       " ['LAFB/LPFB', 'AMI', 'STTC', 'IMI'],\n",
       " ['LAFB/LPFB'],\n",
       " ['AMI', 'STTC'],\n",
       " ['LVH', 'ISCA', 'AMI', 'STTC'],\n",
       " ['NORM'],\n",
       " ['IVCD', 'ISC_'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH', 'AMI', 'STTC', 'ISC_'],\n",
       " ['ISCA', 'IMI'],\n",
       " ['CRBBB'],\n",
       " ['ISCI'],\n",
       " ['LAFB/LPFB'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NST_'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['WPW'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['CLBBB'],\n",
       " ['NORM'],\n",
       " [],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM', 'NST_'],\n",
       " ['STTC', 'RAO/RAE', 'LAO/LAE', '_AVB'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'IMI'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['NORM'],\n",
       " ['NORM', 'IRBBB'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['AMI'],\n",
       " ['NORM'],\n",
       " ['AMI', 'IRBBB'],\n",
       " ['AMI'],\n",
       " ['LVH', 'IVCD', 'ISC_', 'IMI'],\n",
       " ['IRBBB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['CRBBB'],\n",
       " ['NORM'],\n",
       " ['CRBBB', 'AMI', 'IMI'],\n",
       " ['AMI'],\n",
       " [],\n",
       " ['AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI', 'IMI'],\n",
       " ['NST_'],\n",
       " ['NORM'],\n",
       " ['STTC', 'LAO/LAE'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['IVCD'],\n",
       " ['LAFB/LPFB', 'RVH', 'IRBBB'],\n",
       " ['AMI', 'LAO/LAE', '_AVB', 'IMI'],\n",
       " ['ISCA', 'AMI', 'IMI'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['ISCA', 'ISCI'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['AMI', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['LAFB/LPFB', 'IMI'],\n",
       " ['ISCA', 'AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " [],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['CLBBB', 'STTC'],\n",
       " ['IMI'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM', 'IVCD'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['CRBBB', 'LAO/LAE'],\n",
       " ['AMI'],\n",
       " ['LVH', 'ISC_', 'IMI'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'AMI'],\n",
       " ['LVH', 'SEHYP'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'AMI'],\n",
       " ['LVH', 'ISCI'],\n",
       " ['STTC'],\n",
       " ['STTC'],\n",
       " ['ISCA', 'AMI'],\n",
       " ['LVH'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'CRBBB', 'RVH'],\n",
       " ['LAFB/LPFB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['IVCD', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'CRBBB', 'AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['NORM'],\n",
       " ['AMI', 'IVCD', 'NST_'],\n",
       " ['NORM'],\n",
       " ['ISCA'],\n",
       " ['LVH', 'STTC'],\n",
       " ['NORM'],\n",
       " ['CLBBB'],\n",
       " ['CLBBB'],\n",
       " ['LVH'],\n",
       " ['STTC', 'NST_', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NST_'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['ISCA', 'ISCI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['ISCA', 'AMI'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['AMI'],\n",
       " ['NORM'],\n",
       " ['LVH', 'AMI', 'ISC_'],\n",
       " ['IMI'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['IMI'],\n",
       " ['IRBBB'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['STTC', '_AVB', 'IMI'],\n",
       " ['STTC'],\n",
       " ['LVH', 'AMI', 'STTC'],\n",
       " ['NORM'],\n",
       " ['AMI', 'STTC'],\n",
       " ['LAFB/LPFB', 'STTC'],\n",
       " ['LVH', 'ISC_', 'IMI'],\n",
       " ['LAFB/LPFB', 'IMI'],\n",
       " ['NORM', 'IVCD'],\n",
       " ['IRBBB'],\n",
       " ['RAO/RAE'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'AMI', 'IMI'],\n",
       " ['LVH'],\n",
       " ['RAO/RAE'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH', '_AVB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['IRBBB'],\n",
       " ['AMI'],\n",
       " ['AMI'],\n",
       " ['IRBBB'],\n",
       " ['NORM'],\n",
       " ['IRBBB'],\n",
       " ['STTC'],\n",
       " ['LAFB/LPFB', 'STTC', '_AVB'],\n",
       " ['LVH', 'IMI'],\n",
       " ['LVH', 'AMI', 'ISC_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['AMI', 'IMI'],\n",
       " ['NORM'],\n",
       " ['LVH', 'ISC_', 'IMI'],\n",
       " ['AMI', 'IMI', 'RAO/RAE'],\n",
       " ['STTC'],\n",
       " ['AMI', 'IMI'],\n",
       " ['LAFB/LPFB', 'AMI', 'ISCI'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['LVH', 'IMI'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['LVH', 'AMI', 'ISC_'],\n",
       " ['ISCA', 'AMI', '_AVB'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['LAFB/LPFB', 'STTC'],\n",
       " ['ISCA', 'ISCI'],\n",
       " ['ISCA'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NST_'],\n",
       " ['_AVB', 'ISCI'],\n",
       " ['NORM'],\n",
       " ['IRBBB'],\n",
       " ['_AVB'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " [],\n",
       " ['LVH', 'ISCI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH'],\n",
       " ['NORM'],\n",
       " ['ISCA'],\n",
       " ['NORM'],\n",
       " ['NORM', 'IRBBB'],\n",
       " ['ISCA'],\n",
       " ['NORM'],\n",
       " [],\n",
       " ['AMI', 'STTC'],\n",
       " ['NORM'],\n",
       " ['LVH', 'ISC_', 'IMI'],\n",
       " ['NORM'],\n",
       " ['ISCA'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC', 'IRBBB', 'LAO/LAE', '_AVB'],\n",
       " ['LVH', 'IRBBB', '_AVB', 'ISC_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI', 'LAO/LAE', 'IRBBB'],\n",
       " ['LAFB/LPFB', 'AMI'],\n",
       " ['LVH'],\n",
       " ['LVH'],\n",
       " ['ISCA', 'IMI'],\n",
       " ['LVH', 'IRBBB', 'ISC_', 'IMI'],\n",
       " ['IMI'],\n",
       " ['IMI'],\n",
       " ['AMI'],\n",
       " ['STTC'],\n",
       " ['LAFB/LPFB', 'CRBBB', '_AVB', 'RVH'],\n",
       " ['NORM'],\n",
       " [],\n",
       " ['LAFB/LPFB', 'AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'CRBBB', '_AVB'],\n",
       " ['LVH'],\n",
       " ['IRBBB', 'IMI'],\n",
       " ['NORM'],\n",
       " ['LVH', 'AMI', 'ISC_', 'IMI'],\n",
       " ['IRBBB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH', 'ISC_', 'IMI'],\n",
       " ['LAFB/LPFB'],\n",
       " ['LVH', 'LAFB/LPFB', 'ISC_'],\n",
       " ['LVH', 'LAO/LAE', 'LMI', 'ISC_'],\n",
       " ['IVCD', 'IMI'],\n",
       " ['NORM'],\n",
       " ['STTC', 'IMI'],\n",
       " ['LAFB/LPFB', 'CRBBB', 'RVH'],\n",
       " [],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['LVH'],\n",
       " ['LVH', 'CLBBB', 'ISC_'],\n",
       " ['LAFB/LPFB', 'ISCA'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['LVH', 'LAFB/LPFB', 'ISC_'],\n",
       " ['CLBBB'],\n",
       " ['STTC'],\n",
       " ['IMI'],\n",
       " ['STTC'],\n",
       " ['STTC', 'IMI'],\n",
       " ['STTC'],\n",
       " ['CRBBB', 'AMI'],\n",
       " ['LVH', 'IRBBB', 'ISC_'],\n",
       " ['LAFB/LPFB'],\n",
       " ['NORM'],\n",
       " ['LVH'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['LVH'],\n",
       " ['LAFB/LPFB', 'AMI'],\n",
       " ['NORM'],\n",
       " ['CRBBB', 'IMI'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'ISCA'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI', 'IRBBB', 'IMI'],\n",
       " ['IRBBB'],\n",
       " ['IMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['STTC'],\n",
       " [],\n",
       " ['LAFB/LPFB', 'AMI'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['ISC_', 'LVH', 'AMI', 'NST_', 'IMI'],\n",
       " ['ISCA'],\n",
       " ['LVH', 'LAFB/LPFB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['AMI'],\n",
       " ['LVH', 'LAO/LAE', 'NST_', 'ISC_'],\n",
       " ['NORM', 'IVCD'],\n",
       " [],\n",
       " ['NORM'],\n",
       " ['LAFB/LPFB', 'CRBBB', '_AVB', 'IMI'],\n",
       " ['AMI', 'IRBBB'],\n",
       " ['LVH'],\n",
       " ['IMI'],\n",
       " ['LVH', 'AMI', 'LAFB/LPFB', 'ISC_'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['ISCA', 'IMI'],\n",
       " ['LAFB/LPFB', 'AMI'],\n",
       " ['ISCA'],\n",
       " ['NST_'],\n",
       " [],\n",
       " ['LAFB/LPFB', 'IMI'],\n",
       " ['LVH', 'NST_', 'ISC_'],\n",
       " ['IVCD'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['LVH', 'ISC_'],\n",
       " ['CRBBB'],\n",
       " ['NORM'],\n",
       " ['IMI'],\n",
       " ['LVH'],\n",
       " ['NST_'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['STTC'],\n",
       " ['AMI', 'IMI'],\n",
       " ['IRBBB'],\n",
       " ['NST_'],\n",
       " ['LAFB/LPFB'],\n",
       " ['NORM'],\n",
       " ['NORM'],\n",
       " ['CRBBB', 'AMI'],\n",
       " ['NORM'],\n",
       " ...]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c222c806",
   "metadata": {},
   "outputs": [],
   "source": [
    "# leads I \n",
    "list_train= []\n",
    "list_test = []\n",
    "for i in X_train:\n",
    "    t=np.delete(i,lead_types[lead_name], 1)\n",
    "    list_train.append(t)\n",
    "    \n",
    "for i in X_test:\n",
    "    t=np.delete(i,lead_types[lead_name], 1)\n",
    "    list_test.append(t)\n",
    "\n",
    "\n",
    "np.save('x_train.npy',np.array(list_train))\n",
    "np.save('x_test.npy',np.array(list_test))\n",
    "\n",
    "# np.save('y_train.npy',np.array(y_train))\n",
    "# np.save('y_test.npy',np.array(y_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5ea0fd63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pickle y_train and y_test\n",
    "import pickle\n",
    "with open('y_train.pkl', 'wb') as f:\n",
    "    pickle.dump(y_train, f)\n",
    "with open('y_test.pkl', 'wb') as f:\n",
    "    pickle.dump(y_test, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "290f209d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0                         [NORM]\n",
       "1                         [NORM]\n",
       "2                         [NORM]\n",
       "3                         [NORM]\n",
       "4                         [NORM]\n",
       "                  ...           \n",
       "2198                      [NORM]\n",
       "2199    [LAFB/LPFB, CRBBB, _AVB]\n",
       "2200                      [NORM]\n",
       "2201          [LAFB/LPFB, IRBBB]\n",
       "2202                       [IMI]\n",
       "Length: 2203, dtype: object"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_test = pd.Series(y_test)\n",
    "y_test"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ECG_39",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
